import socket
import json
import struct
import numpy as np
from text2vec import SentenceModel
from sklearn.metrics.pairwise import cosine_similarity
import time


# 加载关键词的函数
def load_keywords(path):
    with open(path, 'r', encoding='utf-8') as file:
        # keywords = [line.split()[0] for line in file.readlines()]
        keywords = [line.strip() for line in file.readlines()]
    return keywords


# 初始化 SentenceModel 和加载关键词
model_path = "model/shibing624-text2vec-base-chinese"
model = SentenceModel(model_path)
keywords_file_path = "dict/merged_dictionary.dat"
keywords_dict = load_keywords(keywords_file_path)
keyword_embeddings = model.encode(keywords_dict)


def find_similar_keywords(query, top_n=10):
    query_embedding = model.encode([query])
    similarities = cosine_similarity(query_embedding, keyword_embeddings)
    top_indices = np.argsort(similarities[0])[::-1][:top_n]
    similar_keywords = [keywords_dict[i] for i in top_indices]
    return similar_keywords


def receive_message(conn):
    data = b''
    try:
        while True:
            part = conn.recv(1)
            if not part:
                break  # 客户端已关闭连接
            data += part
            if data.endswith(b'\n'):
                break
    except ConnectionResetError:
        print("Connection was reset by the client.")
    except socket.error as e:
        print(f"Socket error: {e}")
    return data.strip()


def send_message(conn, message):
    try:
        encoded_message = message.encode()
        message_length = len(encoded_message)
        conn.sendall(struct.pack('Q', message_length))
        conn.sendall(encoded_message)
    except ConnectionResetError:
        print("Connection was reset by the client while sending.")
    except socket.error as e:
        print(f"Error sending data: {e}")


def handle_query(data):
    try:
        request = json.loads(data.decode())
        if request["query_id"] == 1:
            query = request["msg"]
            start_time = time.time()
            similar_keywords = find_similar_keywords(query)
            elapsed_time = time.time() - start_time
            response = {
                "msgID": 1,
                "msg": similar_keywords,
                "time": f"{elapsed_time:.4f} seconds"
            }
        else:
            response = {"msgID": 0, "error": "Invalid query_id"}
        return json.dumps(response)
    except json.JSONDecodeError:
        return json.dumps({"msgID": 0, "error": "Invalid JSON format"})


def start_server(host='127.0.0.1', port=8888):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind((host, port))
        s.listen()
        print(f"Server started on {host}:{port}")
        while True:
            try:
                conn, addr = s.accept()
                with conn:
                    print(f"Connected by {addr}")
                    while True:
                        data = receive_message(conn)
                        if not data:
                            break
                        response = handle_query(data)
                        send_message(conn, response)
            except socket.error as e:
                print(f"Server socket error: {e}")


if __name__ == "__main__":
    start_server()
